package dev.alm.langchain4jspringbootdemo.controller;

import dev.alm.langchain4jspringbootdemo.common.ChatRequest;
import dev.alm.langchain4jspringbootdemo.config.AiConfig;
import dev.alm.langchain4jspringbootdemo.config.ChatStore;
import dev.alm.langchain4jspringbootdemo.config.RagConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("ai")
@CrossOrigin
public class ChatController {

    @Resource
    private QwenChatModel qwenChatModel;

    @Resource
    private QwenStreamingChatModel qwenStreamingChatModel;

    @Resource
    private AiConfig.AiAssistant aiAssistant;

    @Resource
    private AiConfig.ChatService chatService;

    @Resource
    private AiConfig.ChatServiceIdStore chatServiceIdStore;

    @Resource
    private RagConfig.RagDemoService ragDemoService;

    @PostMapping("chat")
    public String chat(@RequestBody ChatRequest chatRequest) {
        String question = chatRequest.getQuestion();
        String chat = qwenChatModel.chat(question);
        return chat;
    }

    @RequestMapping("assistant")
    public String assistant(@RequestParam("question") String question) {
        String chat = aiAssistant.chat(question);
        return chat;
    }

    @RequestMapping("chatServiceWithId")
    public String assistantWithId(@RequestParam("memoryId") String memoryId, @RequestParam("question") String question) {
        String chat = chatService.chat(memoryId, question);
        return chat;
    }

    @RequestMapping(value = "/rag", produces = "text/stream;charset=utf-8")
    public Flux<String> rag(@RequestParam("memoryId") String memoryId, @RequestParam("question") String question) {
        return Flux.create(fluxSink -> {
            TokenStream stream = ragDemoService.stream(memoryId, question);
            stream.onPartialResponse(s -> {
                fluxSink.next(s);
            });
            stream.onCompleteResponse(chatResponse -> {
                fluxSink.complete();
            });
            stream.onError(throwable -> {
                fluxSink.error(throwable);
            });
            stream.start();
        });
    }


    @RequestMapping(value = "/streamWithId", produces = "text/stream;charset=utf-8")
    public Flux<String> assistantWithIdStream(@RequestParam("memoryId") String memoryId, @RequestParam("question") String question) {
        System.out.println("question: " + question);
        if (question == null || question.isEmpty()) {
            throw new IllegalArgumentException("Question parameter cannot be empty");
        }
        return Flux.create(fluxSink -> {
            // 动态获取今天的日期
            TokenStream stream = chatService.stream(memoryId, question);
            stream.onPartialResponse(s -> {
                fluxSink.next(s);
            });
            stream.onCompleteResponse(chatResponse -> {
                fluxSink.complete();
            });
            stream.onError(throwable -> {
                fluxSink.error(throwable);
            });
            stream.start();
        });
    }

    @RequestMapping(value = "/as/stream", produces = "text/stream;charset=utf-8")
    public Flux<String> assistantStream(@RequestParam("question") String question) {
        System.out.println("question: " + question);
        if (question == null || question.isEmpty()) {
            throw new IllegalArgumentException("Question parameter cannot be empty");
        }
        return Flux.create(fluxSink -> {
            TokenStream stream = aiAssistant.stream(question);
            stream.onPartialResponse(s -> {
                fluxSink.next(s);
            });
            stream.onCompleteResponse(chatResponse -> {
                fluxSink.complete();
            });
            stream.onError(throwable -> {
                fluxSink.error(throwable);
            });
            stream.start();
        });
    }

    @RequestMapping(value = "chatStore")
    public String chatStore(@RequestParam("memoryId") String memoryId, @RequestParam("question") String question) {
        String chat = chatServiceIdStore.chat(memoryId, question);
        return chat;
    }

    // 解决乱码问题
    @GetMapping(value = "stream", produces = "text/stream;charset=utf-8")
    public Flux<String> stream(@RequestParam("question") String question) {
        Flux<String> flux = Flux.create(fluxSink -> {
            qwenStreamingChatModel.chat(question, new StreamingChatResponseHandler() {

                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });
        });
        return flux;
    }

}
